#coding=utf-8
'''
简单的光线追踪，只考虑反射
'''
import matplotlib.pyplot as plt
import pyopencl as cl
import pyopencl.array as cl_array
import numpy as np
import time
from math import ceil
import os
from os import path
import time
os.environ['PYOPENCL_COMPILER_OUTPUT'] = '0'
# Randomly generate a number between [0, x)
rnd = lambda x: x*np.random.rand()
# 实体的数量
LIGHTS=1
INF=2e10
#7680*1159
GX=1200
GY=800

def rng():
 return np.random.uniform(0,1,1)[0]
'''读取模型数据'''

def flaten(node,coun,bvh):
 coun[0]+=1
 current=coun[0]
 if node.inside: #内部节点，有两个子树
  n0=flaten(node.left,coun,bvh) #左
  n1=flaten(node.right,coun,bvh) #右
  box=node.box.reshape((6,)).tolist()
  bvh[current]=box+[n1,0]  
 else:
  box=node.box.reshape((6,)).tolist()
  bvh[current]=box+[node.offset,node.nPrimitives]

def loadmodel():
 f=open("dragon2.obj",'r')
 ls=f.readlines()
 #print(len(ls))
 f.close()
 points=[]
 normals=[]
 faces=[]
 for line in ls:
     if line[:2]=='v ':
         v=[float(i) for i in line[2:].split()]
         points.append(v)
     elif line[:2]=='vn':
         v=[float(i) for i in line[2:].split()]
         normals.append(v)
     elif line[:2]=='f ':
         '''f 13845/13845/11017 1799/1799/11018 17864/17864/11019 '''
         f=[int(i) for i in line[2:].replace('/',' ').split()]
         p=[points[f[i*3]-1] for i in range(3)]
         n=[normals[f[i*3+2]-1] for i in range(3)]
         faces.append(p+n)
 print('points=',len(points))
 print('normals=',len(normals))
 print('faces=',len(faces))
 faces=np.array(faces)
 info=[]
 for i in range(len(faces)):
  face=faces[i]
  _min=np.array([np.min(faces[i][0:3,0]),np.min(faces[i][0:3,1]),np.min(faces[i][0:3,2])])
  _max=np.array([np.max(faces[i][0:3,0]),np.max(faces[i][0:3,1]),np.max(faces[i][0:3,2])])
  bound=[i,np.array([_min,_max])]
  info.append(bound)    
 start=0
 end=len(info)
 oderinfo=[]
 counter=[0]
 n=NODE(info,start,end,oderinfo,counter,faces)
 oderinfo=np.array(oderinfo,dtype=np.float32).reshape((len(oderinfo),18))
 return len(oderinfo),n.getdata(counter=[0]),oderinfo

class NODE(object):
 left=[]
 right=[]
 box=[]
 inside=True
 nPrimitives=0
 offset=0
 def __init__(self,info,start,end,oderinfo,counter,faces):
  counter[0]+=1
  #排序,根据最大延展方向
  tb=info[start][1]
  for i in range(start,end):
   tb=self.union(tb,info[i][1])
  #找到最大方向
  dim=(tb[1]-tb[0]).argmax()
  #排序
  info[start:end]=sorted(info[start:end],key=lambda s:(np.sum(s[1],axis=0))[dim])
  offset=len(oderinfo);
  if len(info[start:end])==1:
   self.left=''
   self.right='';
   self.inside=False
   self.nPrimitives=1
   tmp=info[start][1]
   for i in range(start,end):
    num,bound=info[i]
    oderinfo.append(faces[num])
    self.offset=offset
    tmp=self.union(tmp,bound)
   self.box=tmp
  else:
   self.left=NODE(info,start,start+(end-start)//2,oderinfo,counter,faces)
   self.right=NODE(info,start+(end-start)//2,end,oderinfo,counter,faces)
  self.box=self.getbox()
   
  def getbox(self):
   if self.inside==False:
    return self.box
   return self.union(self.left.box,self.right.box)
  def union(self,a,b):
   a1=np.array([a[0],b[0]])
   a2=np.array([a[1],b[1]])
   c=np.array([[np.min(a1[:,0]),np.min(a1[:,1]),np.min(a1[:,2])],
                     [np.max(a2[:,0]),np.max(a2[:,1]),np.max(a2[:,2])]])
   return c
  def getdata(self,num):
   coun=[0]
   bvh=[[] for i in range(num)]   
   flaten(self,coun,bvh)
   return np.array(bvh,dtype=np.float32)
   
class RayTracing(object):
 
 def __init__(self):
  '''初始化opencl环境'''
  platforms = cl.get_platforms() #获取可用的opencl平台
  print("%d 个可用的opencl平台"%(len(platforms)))
  for pf in platforms:
   print(pf.get_info(cl.platform_info.NAME),pf.get_info(cl.platform_info.VERSION))
  self.ctx = cl.Context(dev_type=cl.device_type.GPU,properties=[(cl.context_properties.PLATFORM, platforms[1])])
  print("当前使用平台：",platforms[1].get_info(cl.platform_info.NAME))
  self.queue=cl.CommandQueue(self.ctx) #opencl命令队列
  
  # #####################  场景 #####################################
  
  #self.ENTITYS=self.testscene2()
  
  #self.ENTITYS=self.randomscene_with_plate2()
  #NORMAL_UP_X,NORMAL_UP_Y,NORMAL_UP_Z=0,0,1
  #LOOK_AT_X,LOOK_AT_Y,LOOK_AT_Z=0,23,-20
  #LOOK_TO_X,LOOK_TO_Y,LOOK_TO_Z=0,0,-20
  self.ENTITYS=1 #,bvh,faces=loadmodel()
  NORMAL_UP_X,NORMAL_UP_Y,NORMAL_UP_Z=0,1,0
  LOOK_AT_X,LOOK_AT_Y,LOOK_AT_Z=12,2,3
  LOOK_TO_X,LOOK_TO_Y,LOOK_TO_Z=0,0,0  
  # ##################### 定义灯光 ###########位置 颜色 #############
  self.mLights=np.array([5., 5., -10.,0,1.0,1.0,0.0,0.0],dtype=np.float32)  
  #编译内核程序
  opt='-cl-fast-relaxed-math' #'-cl-std=CL2.0' # 
  f=open(path.dirname(__file__)+'\\'+'ray_漫反射_镜面反射_折射.c',encoding='UTF-8')
  src=f.read() #
  f.close()
  # 替换关键字
  src=src.replace('ENTITYS_NUM',str(self.ENTITYS))
  src=src.replace('DIM_X',str(GX))
  src=src.replace('DIM_Y',str(GY))
  src=src.replace('NORMAL_UP_X',str(NORMAL_UP_X))
  src=src.replace('NORMAL_UP_Y',str(NORMAL_UP_Y))
  src=src.replace('NORMAL_UP_Z',str(NORMAL_UP_Z))
  src=src.replace('LOOK_AT_X',str(LOOK_AT_X))
  src=src.replace('LOOK_AT_Y',str(LOOK_AT_Y))
  src=src.replace('LOOK_AT_Z',str(LOOK_AT_Z))
  src=src.replace('LOOK_TO_X',str(LOOK_TO_X))
  src=src.replace('LOOK_TO_Y',str(LOOK_TO_Y))
  src=src.replace('LOOK_TO_Z',str(LOOK_TO_Z))  
  self.prg=cl.Program(self.ctx, src).build(opt) #src:内核源代码
  self.ray_tracing=self.prg.ray_tracing
 
  self.mScene	=np.zeros([GY,GX,4],dtype=np.uint8)
  mf = cl.mem_flags
  ctx=self.ctx
  self.d_mBVH= cl.Buffer(ctx, mf.READ_ONLY | mf.COPY_HOST_PTR, hostbuf=bvh)  
  self.d_mFaces= cl.Buffer(ctx, mf.READ_ONLY | mf.COPY_HOST_PTR, hostbuf=faces)  
  self.d_mLights= cl.Buffer(ctx, mf.READ_ONLY | mf.COPY_HOST_PTR, hostbuf=self.mLights)
  #  #image
  #imf=cl.ImageFormat(cl.channel_order.RGBA,cl.channel_type.UNSIGNED_INT8)
  #self.d_mScene=cl.Image(ctx, cl.mem_flags.WRITE_ONLY, imf, shape=(GX,GY))  
  self.d_mScene=cl.Buffer(ctx, mf.WRITE_ONLY | mf.COPY_HOST_PTR, hostbuf=self.mScene)
  
  self.debugarr=np.zeros(32,dtype=np.float32)
  self.d_debugarr=cl.Buffer(ctx, mf.WRITE_ONLY | mf.COPY_HOST_PTR, hostbuf=self.debugarr)
  #cl.enqueue_nd_range_kernel(self.queue, self.k_initRNG,(config.PARTICLE_NUM,),(64,))
  self.ray_tracing.set_args(self.d_mFaces,self.d_mBVH,self.d_mLights,self.d_mScene,self.d_debugarr)
 def testscene(self):
  # ##################### 定义颜色 ######################
  self.mColor=np.array([1,1,1,0,
                   0.8,0.8,0.0,0,
                   0.8,0.6,0.2,0,
                   0.4,0.8,0.2,0
                   ],dtype=np.float32)
  # ##################### 定义位置 ######################
  self.mPosition=np.array([0.0, 0.0, -1.0, 0.5,
                      0.0, -100.5, -1.0, 100.0,
                      1.0, 0.0, -1.0,0.5,
                      -1,-0.0,-1.0,0.5
                   ],dtype=np.float32)
  
  # ##################### 定义法向量 #####################
  self.mNormal=np.array([0, 0, 0,0,
                    0, 0, 0,0,
                    0, 0, 0,0,
                    0, 0, 0,0
                   ],dtype=np.float32)
  # ##################### 定义材质 ###################### 
  #ambient = .05
  #diffuse_c = 1.
  #specular_c =  1.
  #
  # 第一个0表示该材料镜面反光为0,1表示100%镜面反光，第四项空气折射率/玻璃折射率  1.0/1.5, 如果是inf代表不透明
  self.mMaterial=np.array([1,  1.,  1., 1.5, #第一个球透明，其余都是不透明的
                      0,  1.,  1., INF,
                      1,  0.75,  0.5, INF,
                      0,  0.3,  0.2, INF,
                   ],dtype=np.float32)  
  return 4
 def testscene2(self):
  # ##################### 定义颜色 ######################
  self.mColor=np.array([0.2,0.2,0.2,0,
                   1.0,1.0,1.0,0,
                   0.9,0.76,0.46,0,
                   0.0,0.0,1,0,
                   0.90,0.90,0.90,0
                   ],dtype=np.float32)
  # ##################### 定义位置 ######################
  self.mPosition=np.array([0.0, -10004.0, -20.0, 10000,
                      0.0, -0.0, -20.0, 4.0,
                      5.0, -1, -15,2,
                      5,-0.0,-25.0,3,
                      -5.5,-0.0,-15.0,3
                   ],dtype=np.float32)
  
  # ##################### 定义法向量 #####################
  self.mNormal=np.array([0, 0, 0,0,
                    0, 0, 0,0,
                    0, 0, 0,0,
                    0, 0, 0,0,
                    0, 0, 0,0
                   ],dtype=np.float32)
  # ##################### 定义材质 ###################### 
  #ambient = .05
  #diffuse_c = 1.
  #specular_c =  1.
  #
  # 第一个0表示该材料镜面反光为0,1表示100%镜面反光，第四项空气折射率/玻璃折射率  1.0/1.5, 如果是inf代表不透明
  self.mMaterial=np.array([0,  1.,  1., INF, #第一个球透明，其余都是不透明的
                      1,  1.,  1., 1.6,
                      1,  0.75,  0.5, INF,
                      1,  0.3,  0.2, INF,
                      1,  0.3,  0.2, INF,
                   ],dtype=np.float32)  
  return 5 
 
 def randomscene(self):
  n=500
  self.mColor=np.zeros((n,4),dtype=np.float32)
  self.mPosition=np.zeros((n,4),dtype=np.float32)
  self.mMaterial=np.zeros((n,4),dtype=np.float32)
  self.mNormal=np.zeros((n,4),dtype=np.float32)
  i=0
  # 一个非常大的球
  self.mColor[i,0:3]=[0.5, 0.5, 0.5]
  self.mPosition[i,0:4]=[0., -1000., 0.,1000]
  self.mMaterial[i,0:4]=[0., 0., 0.,INF]
  i+=1
  for a in range(-11,11):
   for b in range(-11,11):
    mat=rng()
    center=np.array([a+0.9*rng(),0.2,b+0.9*rng()])
    if np.linalg.norm(center-[4,0.2,0])>0.9:
     # 散射
     if mat<0.8:
      self.mColor[i,0:3]=[rng(), rng(), rng()]
      self.mPosition[i,0:4]=[center[0], center[1], center[2],0.2]
      self.mMaterial[i,0:4]=[0., 0., 0.,INF]
      i+=1
     # 镜面
     elif mat<0.95:
      self.mColor[i,0:3]=np.array([1+rng(), 1+rng(), 1+rng()])*0.5
      self.mPosition[i,0:4]=[center[0], center[1], center[2],0.2]
      self.mMaterial[i,0:4]=[1, 0., 0.,INF]
      i+=1
     # 透明
     else:
      self.mColor[i,0:3]=[1,1,1]
      self.mPosition[i,0:4]=[center[0], center[1], center[2],0.2]
      self.mMaterial[i,0:4]=[1, 0., 0.,1.5] 
      i+=1
  # 透明大球
  self.mColor[i,0:3]=[1,1,1]
  self.mPosition[i,0:4]=[0, 1, 0,1]
  self.mMaterial[i,0:4]=[1, 0., 0.,1.5] 
  i+=1
  # 哑光大球
  self.mColor[i,0:3]=[0.4,0.4,0.1]
  self.mPosition[i,0:4]=[-4, 1, 0,1]
  self.mMaterial[i,0:4]=[0, 0., 0.,INF] 
  i+=1   
  # 金属大球
  self.mColor[i,0:3]=[0.7,0.6,0.5]
  self.mPosition[i,0:4]=[4, 1, 0,1]
  self.mMaterial[i,0:4]=[1, 0., 0.,INF] 
  i+=1
  print('i= ===',i)
  return i
 def randomscene_with_plate(self):
  n=500
  self.mColor=np.zeros((n,4),dtype=np.float32)
  self.mPosition=np.zeros((n,4),dtype=np.float32)
  self.mMaterial=np.zeros((n,4),dtype=np.float32)
  self.mNormal=np.zeros((n,4),dtype=np.float32)
  i=0
  # 一个非常大的球
  #self.mColor[i,0:3]=[0.5, 0.5, 0.5]
  #self.mPosition[i,0:4]=[0., -1000., 0.,1000]
  #self.mMaterial[i,0:4]=[0., 0., 0.,INF]
  #i+=1
  for a in range(-11,11):
   for b in range(-11,11):
    mat=rng()
    center=np.array([a+0.9*rng(),0.2,b+0.9*rng()])
    if np.linalg.norm(center-[4,0.2,0])>0.9:
     # 散射
     if mat<0.8:
      self.mColor[i,0:3]=[rng(), rng(), rng()]
      self.mPosition[i,0:4]=[center[0], center[1], center[2],0.2]
      self.mMaterial[i,0:4]=[0., 0., 0.,INF]
      i+=1
     # 镜面
     elif mat<0.95:
      self.mColor[i,0:3]=np.array([1+rng(), 1+rng(), 1+rng()])*0.5
      self.mPosition[i,0:4]=[center[0], center[1], center[2],0.2]
      self.mMaterial[i,0:4]=[1, 0., 0.,INF]
      i+=1
     # 透明
     else:
      self.mColor[i,0:3]=[1,1,1]
      self.mPosition[i,0:4]=[center[0], center[1], center[2],0.2]
      self.mMaterial[i,0:4]=[1, 0., 0.,1.5] 
      i+=1
  # 透明大球
  self.mColor[i,0:3]=[1,1,1]
  self.mPosition[i,0:4]=[0, 1, 0,1]
  self.mMaterial[i,0:4]=[1, 0., 0.,1.5] 
  i+=1
  # 哑光大球
  self.mColor[i,0:3]=[0.4,0.4,0.1]
  self.mPosition[i,0:4]=[-4, 1, 0,1]
  self.mMaterial[i,0:4]=[0, 0., 0.,INF] 
  i+=1   
  # 金属大球
  self.mColor[i,0:3]=[0.7,0.6,0.5]
  self.mPosition[i,0:4]=[4, 1, 0,1]
  self.mMaterial[i,0:4]=[1, 0., 0.,INF] 
  i+=1
  # 平板
  self.mColor[i,0:3]=[0.7,0.6,0.5]
  self.mPosition[i,0:4]=[0, -0.1, 0,INF]
  self.mMaterial[i,0:4]=[0, 0., 0.,INF]   
  self.mNormal[i]=[0,1,0,0]
  i+=1
  print('i= ===',i)
  return i 
 def randomscene_with_plate2(self):
  n=2
  self.mColor=np.zeros((n,4),dtype=np.float32)
  self.mPosition=np.zeros((n,4),dtype=np.float32)
  self.mMaterial=np.zeros((n,4),dtype=np.float32)
  self.mNormal=np.zeros((n,4),dtype=np.float32)
  i=0
  # 透明大球
  self.mColor[i,0:3]=[1,1,1]
  self.mPosition[i,0:4]=[0, 0, -20,4]
  self.mMaterial[i,0:4]=[1, 0., 0.,1.6] 
  i+=1
  # 平板
  self.mColor[i,0:3]=[0.5,0.2,0.2]
  self.mPosition[i,0:4]=[0, -4.1, -0,INF]
  self.mMaterial[i,0:4]=[1, 0., 0.,INF]   
  self.mNormal[ i]=[0,1,0,0]
  i+=1
  print('i= ===',i)
  return i 
 def run(self):
  cl.enqueue_nd_range_kernel(self.queue, self.ray_tracing,(GX,GY,),(8,8,))
  #cl.enqueue_copy(self.queue, self.mScene, self.d_mScene)
  cl.enqueue_copy(self.queue,self.mScene,self.d_mScene,origin=(0,0),region=(GX,GY))
 def show(self):
  fig, axs = plt.subplots(1, 1)
  #self.debugarr+=0.99999
  
   #bitmap = np.transpose(bitmap/255., (1, 0, 2)) # swap image's x-y axes
  
  axs.imshow(self.mScene)
  axs.grid(False)
 
   # sort the spheres by Z for visualizing z-level order, plot using circle artists
  #for ax in axs:
  axs.set_xlim(0, GX)
  axs.set_ylim(0, GY)
  #print(self.debugarr) 
  #np.savetxt('c:/lines.txt',self.debugarr)
  plt.imsave('fig2.png', np.rot90(self.mScene,k=2))
  plt.tight_layout()
  plt.show()  
 def savetovtk(self):
  cl.enqueue_copy(self.queue, self.debugarr, self.d_debugarr)
  vtk=open("LINES.vtk",'w')
  print("# vtk DataFile Version 1.0",file=vtk)
  print("Line representation of vtk",file=vtk)
  print("ASCII",file=vtk)
  print("",file=vtk)
  print("DATASET POLYDATA",file=vtk)
  arr=self.debugarr.reshape((1024,4))
  num=min(512,int(arr[0,3]))
  print("POINTS %d float"%(2*num),file=vtk)
  for i in range(2*num):
   print("%f %f %f"%(arr[i,0],arr[i,1],arr[i,2]),file=vtk)
  print('',file=vtk)
  print("LINES %d %d"%(2*num,num*3),file=vtk)
  for i in range(num):
   print('2 %d %d'%(i*2,i*2+1),file=vtk)
  vtk.close()
 
if __name__ == '__main__':

 p=RayTracing()
 t0=time.time()
 #p.randomscene()
 p.run()
 print('time clapsed=',time.time()-t0)
 #p.savetovtk()
 p.show()